package com.example.rocket.config;

import com.example.rocket.socket.MyWebSocketHandler;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @author jwz1106
 */
@Configuration
@EnableWebSocket // 用于开启注解接收和发送消息
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(new MyWebSocketHandler(), "ws") //设置处理类和连接路径
            .setAllowedOrigins("*")//设置作用域
            .addInterceptors(new MyWebSocketInterceptor());
    }

    class MyWebSocketInterceptor implements HandshakeInterceptor {

        @Override
        public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
            System.out.println("前置拦截~~");
            HttpHeaders headers = request.getHeaders();
            List<String> tokens = headers.get("token");
            if (Objects.isNull(tokens) || tokens.isEmpty()) {
                return false;
            }
            String token = tokens.get(0);
            attributes.put("username", token);
            return true;
        }

        @Override
        public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
            System.out.println("后置拦截~~");
        }
    }
}